/**
 * @file   Twogrid.h
 * @author ubuntu <dadi@ubuntu>
 * @date   Tue May 31 13:07:23 2022
 * 
 * @brief  Two Grid Scheme.
 * 
 * 
 */

#ifndef __DADI_TWOGRID__
#define __DADI_TWOGRID__

#include <iostream>
#include <vector>
#include <string>
#include <cmath>
#include <valarray>
#include <lapacke.h>

class TwoGrid
{
protected:
    std::valarray<double> Originalf;
    std::vector<std::valarray<double>> v;	/**< Numerical solution of each level. */
    std::vector<std::valarray<double>> f;	/**< RHS of each level. */
    std::string relaxation_scheme; /**< Relaxation_scheme, Weighted-Jacobi, SOR. */
    double w;			/**< Weight coefficient for relaxation scheme. */
    double accuracy;		/**< One of the iteration termination conditions: the accuracy of the solution. */
    int max_iteration;		/**< One of the iteration termination conditions: the maximum number of iterations. */
    
public:
    void Set_relaxation_scheme(std::string _relaxation_scheme, double _w);
    void Set_termination_conditions(double _accuracy, int _max_iteration);
    std::valarray<double> Solve();
    
    virtual void Initialize(std::valarray<double> _v, double (*_f)(double *x), double (*_g)(double *x), std::vector<std::vector<double>> Region) = 0;
    virtual std::valarray<double> Grid_operator(std::valarray<double> _v) = 0;
    virtual std::valarray<double> Restriction_operator(std::valarray<double> _v) = 0;
    virtual std::valarray<double> Interpolation_operator(std::valarray<double> _v) = 0;
    virtual std::valarray<double> Relaxation_operator(std::valarray<double> _v, std::valarray<double> _f) = 0;
    virtual void Two_grid() = 0;
};

void TwoGrid::Set_relaxation_scheme(std::string _relaxation_scheme, double _w)
{
    relaxation_scheme = _relaxation_scheme;
    w = _w;
};

void TwoGrid::Set_termination_conditions(double _accuracy, int _max_iteration)
{
    accuracy = _accuracy;
    max_iteration = _max_iteration;
};

std::valarray<double> TwoGrid::Solve()
{
    double a = 0.0;
    int k = 0;
    std::valarray<double> r(v[0].size());
    std::valarray<double> zero(0.0, v[0].size());
    std::valarray<double> solution = v[0];
    r = Originalf - Grid_operator(solution);
    f[0] = r;
    v[0] = zero;
    r *= r;
    for(int i = 0; i < v[0].size(); i++)
	a = a + r[i];
    a = sqrt(a);
    
    while(a >= accuracy && k < max_iteration)
    {
	a = 0.0;
	Two_grid();
	solution = solution + v[0];
	r = Originalf - Grid_operator(solution);
	f[0] = r;
	v[0] = zero;
	r *= r;
	for(int i = 0; i < v[0].size(); i++)
	    a = a + r[i];
	a = sqrt(a);
	k++;
	std::cout << "Iteration" << k << ": The L2 norm of the residual is " << a << "." << std::endl;
    }
    return solution;
};

class TwoGrid1D : public TwoGrid
{
private:
    double x0;
    double x1;
public:
    void Initialize(std::valarray<double> _v, double (*_f)(double *x), double (*_g)(double *x), std::vector<std::vector<double>> Region) override;
    std::valarray<double> Grid_operator(std::valarray<double> _v) override;
    std::valarray<double> Restriction_operator(std::valarray<double> _v) override;
    std::valarray<double> Interpolation_operator(std::valarray<double> _v) override;
    std::valarray<double> Relaxation_operator(std::valarray<double> _v, std::valarray<double> _f) override;
    void Two_grid() override;
};

void TwoGrid1D::Initialize(std::valarray<double> _v, double (*_f)(double *x), double (*_g)(double *x), std::vector<std::vector<double>> Region)
{
    v.resize(2);
    f.resize(2);
    v[0] = _v;
    int n = _v.size();
    f[0].resize(n);
    x0 = Region[0][0];
    x1 = Region[1][0];
    double h = (x1-x0)/(n+1);
    double c[1];
    for(int i = 0; i < n; i++)
    {
	c[0] = x0 + (i+1)*h;
	f[0][i] = _f(c);
    }
    //Dirichlet condition.
    c[0] = x0;
    f[0][0] += _g(c)/(h*h);
    c[0] = x1;
    f[0][n-1] += _g(c)/(h*h);
    Originalf = f[0];
};

std::valarray<double> TwoGrid1D::Grid_operator(std::valarray<double> _v)
{
    int n = _v.size();
    double h = (x1-x0)/(n+1);
    std::valarray<double> u(n);
    for(int i = 0; i < n; i++)
    {
	if(i == 0)
	    u[i] = (2.0*_v[i] - 1.0*_v[i+1])/(h*h);
	else if(i == n-1)
	    u[i] = (-1.0*_v[i-1] + 2.0*_v[i])/(h*h);
	else
	    u[i] = (-1.0*_v[i-1] + 2.0*_v[i] - 1.0*_v[i+1])/(h*h);
    }
    return u;
};

std::valarray<double> TwoGrid1D::Restriction_operator(std::valarray<double> _v)
{
    int n = _v.size();
    std::valarray<double> u((n-1)/2);
    for(int i = 1; i < (n+1)/2; i++)
	u[i-1] = (_v[2*i-2] + 2.0*_v[2*i-1] + _v[2*i])/4.0;
    return u;
};

std::valarray<double> TwoGrid1D::Interpolation_operator(std::valarray<double> _v)
{
    int n = _v.size();
    std::valarray<double> u(2*n+1);
    for(int i = 1; i < 2*n+2; i++)
    {
	if(i%2 == 0)
	    u[i-1] = _v[i/2-1];
	else if(i == 1)
	    u[i-1] = _v[(i-1)/2]/2.0;
	else if(i == 2*n+1)
	    u[i-1] = _v[(i-1)/2-1]/2.0;
	else
	    u[i-1] = (_v[(i-1)/2-1] + _v[(i-1)/2])/2.0;
    }
    return u;
};

std::valarray<double> TwoGrid1D::Relaxation_operator(std::valarray<double> _v, std::valarray<double> _f)
{
    int n = _v.size();
    double h = (x1-x0)/(n+1);
    std::valarray<double> u(n);
    if(relaxation_scheme == "Weighted_Jacobi")
    {
	for(int i = 0; i < n; i++)
	{
	    if(i == 0)
		u[i] = -1.0*_v[i+1];
	    else if(i == n-1)
		u[i] = -1.0*_v[i-1];
	    else
		u[i] = -1.0*_v[i-1]-1.0*_v[i+1];
	    u[i] = -(u[i] - _f[i]*h*h)/2.0;
	}
	u = (1-w)*_v + w*u;
    }
    else
    {
	u = _v;
	for(int i = 0; i < n ;i++)
	{
	    if(i == 0)
		u[i] = -1.0*u[i+1];
	    else if(i == n-1)
		u[i] = -1.0*u[i-1];
	    else
		u[i] = -1.0*u[i-1]-1.0*u[i+1];
	    u[i] = -(u[i] - _f[i]*h*h)/2.0;
	}
	u = (1-w)*_v + w*u;
    }
    return u;
};

void TwoGrid1D::Two_grid()
{
    for(int i = 0; i < 2; i++)
	v[0] = Relaxation_operator(v[0], f[0]);
    f[1] = Restriction_operator(f[0]-Grid_operator(v[0]));
    int n = f[1].size();
    v[1].resize(n);

    //directly solve
    double h = (x1-x0)/(n+1);
    double A[n*n] = {0.0};
    double Rhs[n];
    for(int i = 0; i < n; i++)
    {
	if(i == 0)
	{
	    A[i*n+i] = 2.0/(h*h);
	    A[i*n+i+1] = -1.0/(h*h);
	}
	else if(i == n-1)
	{
	    A[i*n+i] = 2.0/(h*h);
	    A[i*n+i-1] = -1.0/(h*h);
	}
	else
	{
	    A[i*n+i] = 2.0/(h*h);
	    A[i*n+i+1] = -1.0/(h*h);
	    A[i*n+i-1] = -1.0/(h*h);
	}
	Rhs[i] = f[1][i];
    }
    //Cholesky.
    LAPACKE_dposv(LAPACK_COL_MAJOR, 'L', n, 1, A, n, Rhs, n);
    for(int i = 0; i < n; i++)
	v[1][i] = Rhs[i];

    v[0] = v[0] + Interpolation_operator(v[1]);
    for(int i = 0; i < 2; i++)
	v[0] = Relaxation_operator(v[0], f[0]);
};

class TwoGrid2D : public TwoGrid
{
private:
    double x0;
    double x1;
    double y0;
    double y1;
public:
    void Initialize(std::valarray<double> _v, double (*_f)(double *x), double (*_g)(double *x), std::vector<std::vector<double>> Region) override;
    std::valarray<double> Grid_operator(std::valarray<double> _v) override;
    std::valarray<double> Restriction_operator(std::valarray<double> _v) override;
    std::valarray<double> Interpolation_operator(std::valarray<double> _v) override;
    std::valarray<double> Relaxation_operator(std::valarray<double> _v, std::valarray<double> _f) override;
    void Two_grid() override;
};

void TwoGrid2D::Initialize(std::valarray<double> _v, double (*_f)(double *x), double (*_g)(double *x), std::vector<std::vector<double>> Region)
{
    v.resize(2);
    f.resize(2);
    v[0] = _v;
    int n = _v.size();
    int m = (int)sqrt(n);
    f[0].resize(n);
    x0 = Region[0][0];
    y0 = Region[0][1];
    x1 = Region[1][0];
    y1 = Region[1][1];
    double hx = (x1-x0)/(m+1);
    double hy = (y1-y0)/(m+1);
    double c[2];
    // int ix;
    // int iy;
    for(int i = 0; i < m; i++)
    {
	for(int j = 0; j < m; j++)
	{
	    c[0] = (j+1)*hx;
	    c[1] = (i+1)*hy;
	    f[0][i*m+j] = _f(c);
	}
    }
    // for(int i = 0; i < n; i++)
    // {
    // 	ix = i%m;
    // 	iy = i/m;
    // 	c[0] = (ix+1)*hx;
    // 	c[1] = (iy+1)*hy;
    // 	f[0][i] = _f(c);
    // }
    //Dirichlet condition.
    for(int i = 0; i < m; i++)
    {
	//Bottom
	c[0] = x0 + (i+1)*hx;
	c[1] = y0;
	f[0][i] += _g(c)/(hx*hx);

	//Left
	c[0] = x0;
	c[1] = y0 + (i+1)*hy;
	f[0][i*m] += _g(c)/(hx*hx);

	//Right
	c[0] = x1;
	c[1] = y0 + (i+1)*hy;
	f[0][(i+1)*m-1] += _g(c)/(hx*hx);

	//Top
	c[0] = x0 + (i+1)*hx;
	c[1] = y1;
	f[0][i+(m-1)*m] += _g(c)/(hx*hx);
    }
    Originalf = f[0];
};

std::valarray<double> TwoGrid2D::Grid_operator(std::valarray<double> _v)
{
    int n = _v.size();
    int m = (int)sqrt(n);
    double h = (x1-x0)/(m+1);
    std::valarray<double> u(n);
    for(int i = 0; i < m; i++)
    {
	if(i == 0)
	{
	    u[i] = (4.0*_v[i] - 1.0*_v[i+1] - 1.0*_v[i+m])/(h*h);
	}
	else if(i == m-1)
	{
	    u[i] = (-1.0*_v[i-1] + 4.0*_v[i] - 1.0*_v[i+m])/(h*h);
	}
	else
	{
	    u[i] = (-1.0*_v[i-1] + 4.0*_v[i] - 1.0*_v[i+1] - 1.0*_v[i+m])/(h*h);
	}
    }
    for(int j = 1; j <= m-2; j++)
    {
	for(int i = j*m; i < (j+1)*m; i++)
	{
	    if(i == j*m)
	    {
		u[i] = (4.0*_v[i] - 1.0*_v[i+1] - 1.0*_v[i+m] - 1.0*_v[i-m])/(h*h);
	    }
	    else if(i == (j+1)*m-1)
	    {
		u[i] = (-1.0*_v[i-1] + 4.0*_v[i] - 1.0*_v[i+m] - 1.0*_v[i-m])/(h*h);
	    }
	    else
	    {
		u[i] = (-1.0*_v[i-1] + 4.0*_v[i] - 1.0*_v[i+1] - 1.0*_v[i+m] - 1.0*_v[i-m])/(h*h);
	    }
	}
    }
    for(int i = m*(m-1); i < n; i++)
    {
	if(i == m*(m-1))
	{
	    u[i] = (4.0*_v[i] - 1.0*_v[i+1] - 1.0*_v[i-m])/(h*h);
	}
	else if(i == n-1)
	{
	    u[i] = (-1.0*_v[i-1] + 4.0*_v[i] - 1.0*_v[i-m])/(h*h);
	}
	else
	{
	    u[i] = (-1.0*_v[i-1] + 4.0*_v[i] - 1.0*_v[i+1] - 1.0*_v[i-m])/(h*h);
	}
    }
    return u;
};

std::valarray<double> TwoGrid2D::Restriction_operator(std::valarray<double> _v)
{
    int n = _v.size();
    int m = (int)sqrt(n);
    int k = (m-1)/2;
    // int ix;
    // int iy;
    int id;
    std::valarray<double> u(k*k);
    for(int i = 0; i < k; i++)
    {
	for(int j = 0; j < k; j++)
	{
	    id = 2*j+1+(2*i+1)*m;
	    u[j+i*k] = (_v[id-m-1] + _v[id-m+1] + _v[id+m-1] + _v[id+m+1] +
		      2.0*(_v[id-1] + _v[id+1] + _v[id-m] + _v[id+m]) + 4.0*_v[id])/16.0;
	}
    }
    // for(int i = 0; i < k*k; i++)
    // {
    // 	ix = i%k;
    // 	iy = i/k;
    // 	id = 2*ix+1+(2*iy+1)*m;
    // 	u[ix+iy*k] = (_v[id-m-1] + _v[id-m+1] + _v[id+m-1] + _v[id+m+1] +
    // 		      2.0*(_v[id-1] + _v[id+1] + _v[id-m] + _v[id+m]) + 4.0*_v[id])/16.0;
    // }
    return u;
};

std::valarray<double> TwoGrid2D::Interpolation_operator(std::valarray<double> _v)
{
    int n = _v.size();
    int m = (int)sqrt(n);
    int k = 2*m+1;
    int id1;
    int id2;
    std::valarray<double> u(k*k);
    for(int i = 1; i <= m-1; i++)
    {
	for(int j = 1; j <= m-1; j++)
	{
	    id1 = j-1+(i-1)*m;
	    id2 = 2*j-1+(2*i-1)*k;
	    u[id2] = _v[id1];
	    u[id2+1] = (_v[id1] + _v[id1+1])/2.0;
	    u[id2+k] = (_v[id1] + _v[id1+m])/2.0;
	    u[id2+k+1] = (_v[id1] + _v[id1+1] + _v[id1+m] + _v[id1+m+1])/4.0;
	}
    }
    for(int j = 0; j <= m; j++)
    {
	if(j == 0)
	{
	    u[0] = _v[0]/4.0;
	    u[k*(k-1)] = _v[m*(m-1)]/4.0;
	    u[k*(k-2)] = _v[m*(m-1)]/2.0;
	}
	else if(j == m)
	{
	    u[k-1] = _v[m-1]/4.0;
	    u[k-2] = _v[m-1]/2.0;
	    u[k*k-1] = _v[m*m-1]/4.0;
	    u[k*k-2] = _v[m*m-1]/2.0;
	    u[k*(k-1)-1] = _v[m*m-1]/2.0;
	    u[k*(k-1)-2] = _v[m*m-1];
	}
	else
	{
	    u[2*j-1] = _v[j-1]/2.0;
	    u[2*j] = (_v[j-1]+_v[j])/4.0;
	    u[2*j+(k-1)*k] = (_v[j-1+(m-1)*m]+_v[j+(m-1)*m])/4.0;
	    u[2*j+(k-1)*k-1] = _v[j-1+(m-1)*m]/2.0;
	    u[2*j+(k-2)*k] = (_v[j-1+(m-1)*m]+_v[j+(m-1)*m])/2.0;
	    u[2*j+(k-2)*k-1] = _v[j-1+(m-1)*m];
	}
    }
    for(int i = 1; i <= m-1; i++)
    {
	u[(2*i-1)*k] = _v[(i-1)*m]/2.0;
	u[2*i*k] = (_v[(i-1)*m] + _v[i*m])/4.0;
	u[2*i*k-1] = _v[i*m-1]/2.0;
	u[2*i*k-2] = _v[i*m-1];
	u[(2*i+1)*k-1] = (_v[i*m-1] + _v[(i+1)*m-1])/4.0;
	u[(2*i+1)*k-2] = (_v[i*m-1] + _v[(i+1)*m-1])/2.0;
    }
    return u;
};

std::valarray<double> TwoGrid2D::Relaxation_operator(std::valarray<double> _v, std::valarray<double> _f)
{
    int n = _v.size();
    int m = (int)sqrt(n);
    double h = (x1-x0)/(m+1);
    std::valarray<double> u(n);
    if(relaxation_scheme == "Weighted_Jacobi")
    {
	u = Grid_operator(u);
	u = -(u - _f)*h*h/4.0 + _v;
	u = (1-w)*_v + w*u;
    }
    else
    {
	u = _v;
	for(int i = 0; i < m; i++)
	{
	    if(i == 0)
	    {
		u[i] = (-1.0*u[i+1] - 1.0*u[i+m])/(h*h);
	    }
	    else if(i == m-1)
	    {
		u[i] = (-1.0*u[i-1] - 1.0*u[i+m])/(h*h);
	    }
	    else
	    {
		u[i] = (-1.0*u[i-1] - 1.0*u[i+1] - 1.0*u[i+m])/(h*h);
	    }
	    u[i] = -(u[i] - _f[i])*h*h/4.0;
	}
	for(int j = 1; j <= m-2; j++)
	{
	    for(int i = j*m; i < (j+1)*m; i++)
	    {
		if(i == j*m)
		{
		    u[i] = (-1.0*u[i+1] - 1.0*u[i+m] - 1.0*u[i-m])/(h*h);
		}
		else if(i == (j+1)*m-1)
		{
		    u[i] = (-1.0*u[i-1] - 1.0*u[i+m] - 1.0*u[i-m])/(h*h);
		}
		else
		{
		    u[i] = (-1.0*u[i-1] - 1.0*u[i+1] - 1.0*u[i+m] - 1.0*u[i-m])/(h*h);
		}
		u[i] = -(u[i] - _f[i])*h*h/4.0;
	    }
	}
	for(int i = m*(m-1); i < n; i++)
	{
	    if(i == m*(m-1))
	    {
		u[i] = (-1.0*u[i+1] - 1.0*u[i-m])/(h*h);
	    }
	    else if(i == n-1)
	    {
		u[i] = (-1.0*u[i-1] - 1.0*u[i-m])/(h*h);
	    }
	    else
	    {
		u[i] = (-1.0*u[i-1] - 1.0*u[i+1] - 1.0*u[i-m])/(h*h);
	    }
	    u[i] = -(u[i] - _f[i])*h*h/4.0;
	}
	u = (1-w)*_v + w*u;
    }
    return u;
};

void TwoGrid2D::Two_grid()
{
    for(int i = 0; i < 2; i++)
	v[0] = Relaxation_operator(v[0], f[0]);
    f[1] = Restriction_operator(f[0]-Grid_operator(v[0]));
    int n = f[1].size();
    v[1].resize(n);

    //directly solve
    int m = (int)sqrt(n);
    double h = (x1-x0)/(n+1);
    double A[n*n] = {0.0};
    double Rhs[n];
    for(int i = 0; i < m; i++)
    {
	if(i == 0)
	{
	    A[i*n+i] = 4.0/(h*h);
	    A[i*n+i+1] = -1.0/(h*h);
	    A[i*n+i+m] = -1.0/(h*h);
	}
	else if(i == m-1)
	{
	    A[i*n+i] = 4.0/(h*h);
	    A[i*n+i-1] = -1.0/(h*h);
	    A[i*n+i+m] = -1.0/(h*h);
	}
	else
	{
	    A[i*n+i] = 4.0/(h*h);
	    A[i*n+i-1] = -1.0/(h*h);
	    A[i*n+i+1] = -1.0/(h*h);
	    A[i*n+i+m] = -1.0/(h*h);
	}
	Rhs[i] = f[1][i];
    }
    for(int j = 1; j <= m-2; j++)
    {
	for(int i = j*m; i < (j+1)*m; i++)
	{
	    if(i == j*m)
	    {
		A[i*n+i] = 4.0/(h*h);
		A[i*n+i+1] = -1.0/(h*h);
		A[i*n+i+m] = -1.0/(h*h);
		A[i*n+i-m] = -1.0/(h*h);
	    }
	    else if(i == (j+1)*m-1)
	    {
		A[i*n+i] = 4.0/(h*h);
		A[i*n+i-1] = -1.0/(h*h);
		A[i*n+i+m] = -1.0/(h*h);
		A[i*n+i-m] = -1.0/(h*h);
	    }
	    else
	    {
		A[i*n+i] = 4.0/(h*h);
		A[i*n+i-1] = -1.0/(h*h);
		A[i*n+i+1] = -1.0/(h*h);
		A[i*n+i+m] = -1.0/(h*h);
		A[i*n+i-m] = -1.0/(h*h);
	    }
	    Rhs[i] = f[1][i];
	}
    }
    for(int i = m*(m-1); i < n; i++)
    {
	if(i == m*(m-1))
	{
	    A[i*n+i] = 4.0/(h*h);
	    A[i*n+i+1] = -1.0/(h*h);
	    A[i*n+i-m] = -1.0/(h*h);
	}
	else if(i == n-1)
	{
	    A[i*n+i] = 4.0/(h*h);
	    A[i*n+i-1] = -1.0/(h*h);
	    A[i*n+i-m] = -1.0/(h*h);
	}
	else
	{
	    A[i*n+i] = 4.0/(h*h);
	    A[i*n+i-1] = -1.0/(h*h);
	    A[i*n+i+1] = -1.0/(h*h);
	    A[i*n+i-m] = -1.0/(h*h);
	}
	Rhs[i] = f[1][i];
    }
    //Cholesky.
    LAPACKE_dposv(LAPACK_COL_MAJOR, 'L', n, 1, A, n, Rhs, n);
    for(int i = 0; i < n; i++)
	v[1][i] = Rhs[i];

    v[0] = v[0] + Interpolation_operator(v[1]);
    for(int i = 0; i < 2; i++)
	v[0] = Relaxation_operator(v[0], f[0]);
};


#else
//Do nothing.
#endif
